"""
颜色识别模块
用于识别图像中的特定颜色目标
"""

import cv2
import numpy as np
from typing import Dict, Any, Optional, Tuple
from .vision_base import VisionBase
from lab_conf import color_range
import time


class ColorDetection(VisionBase):
    """颜色识别类"""

    def __init__(self, config: Dict[str, Any]):
        """
        初始化颜色识别模块

        Args:
            config: 配置字典，包含颜色识别参数
        """
        super().__init__(config)
        self.target_color = config.get("target_color", "red")
        self.min_area = config.get("min_area", 300)
        self.color_list = []
        self.color_history_size = config.get("color_history_size", 10)

    def _get_max_contour(self, contours: list) -> Tuple[Optional[np.ndarray], float]:
        """
        获取面积最大的轮廓

        Args:
            contours: 轮廓列表

        Returns:
            最大轮廓和其面积
        """
        max_area = 0
        max_contour = None

        for contour in contours:
            area = abs(cv2.contourArea(contour))
            if area > max_area and area > self.min_area:
                max_area = area
                max_contour = contour

        return max_contour, max_area

    async def process_frame(self, frame: np.ndarray) -> Dict[str, Any]:
        """
        处理图像帧，识别目标颜色

        Args:
            frame: 输入图像帧

        Returns:
            处理结果字典
        """
        # 图像预处理
        frame = cv2.resize(frame, (self.width, self.height))
        frame = cv2.GaussianBlur(frame, (3, 3), 0)
        lab = cv2.cvtColor(frame, cv2.COLOR_BGR2LAB)

        # 颜色分割
        mask = cv2.inRange(
            lab, color_range[self.target_color][0], color_range[self.target_color][1]
        )

        # 形态学处理
        kernel = np.ones((3, 3), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)

        # 查找轮廓
        contours = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)[-2]
        max_contour, max_area = self._get_max_contour(contours)

        result = {
            "detected": False,
            "color": None,
            "center": None,
            "radius": 0,
            "area": 0,
        }

        if max_contour is not None:
            # 获取最小外接圆
            (center, radius) = cv2.minEnclosingCircle(max_contour)
            center = (int(center[0]), int(center[1]))
            radius = int(radius)

            # 更新颜色历史
            self.color_list.append(1)  # 1表示检测到目标颜色
            if len(self.color_list) > self.color_history_size:
                self.color_list.pop(0)

            # 根据历史判断颜色
            if sum(self.color_list) / len(self.color_list) > 0.7:
                result.update(
                    {
                        "detected": True,
                        "color": self.target_color,
                        "center": center,
                        "radius": radius,
                        "area": max_area,
                    }
                )

            # 在图像上绘制结果
            cv2.circle(frame, center, radius, (0, 255, 0), 2)

        # 添加调试信息
        if self.debug:
            frame = self.add_text_to_image(
                frame, f"Color: {self.target_color}", (10, 30)
            )
            frame = self.add_text_to_image(
                frame, f"Detected: {result['detected']}", (10, 60)
            )

        return {"result": result, "frame": frame}

    async def get_vision_data(self) -> Dict[str, Any]:
        """
        获取视觉处理数据

        Returns:
            包含处理结果和时间戳的字典
        """
        frame = await self.get_frame()
        if frame is not None:
            result = await self.process_frame(frame)
            return {
                **result["result"],
                "timestamp": time.time(),
                "frame": result["frame"] if self.debug else None,
            }
        return {
            "detected": False,
            "color": None,
            "center": None,
            "radius": 0,
            "area": 0,
            "timestamp": time.time(),
            "frame": None,
        }
